#include "pch.h"
#include "AssocStateBase.h"
#include "AsyncLibImpl.h"
#include "AutoRefPtr.h"

namespace AsyncLib
{

CAssocStateBase::CAssocStateBase() : m_status(Pending), m_atid(0), m_bExIgnored(false), m_bExHandled(false)
{
}

CAssocStateBase::~CAssocStateBase()
{
	assert(m_status != Pending);

	if (m_exception && !m_bExHandled)
	{
#ifdef _WIN32
		try
		{
			std::rethrow_exception(m_exception);
		}
		catch ([[maybe_unused]] const std::exception& e)
		{
			_RPTF1(_CRT_WARN, "Unhandled exception: %s\n", e.what());
		}
		catch (...)
		{
			_RPTF0(_CRT_WARN, "Unhandled exception.\n");
		}
#endif
		if (!m_bExIgnored)
			std::abort();
	}
}

CAssocStateBase::Status CAssocStateBase::SetCallback(std::function<void(bool)>&& fnCallback)
{
	std::unique_lock lock(m_mutex);
	assert(!m_fnCallback);

	if (m_status == Pending)
		m_fnCallback = std::move(fnCallback);

	return m_status;
}

bool CAssocStateBase::IsCompleted() const
{
	std::unique_lock lock(m_mutex);
	return m_status == Completed;
}

void CAssocStateBase::Forget()
{
	m_bExIgnored = true;
}

bool CAssocStateBase::Wait()
{
	std::unique_lock lock(m_mutex);

	if (m_status == Pending)
	{
		std::condition_variable condVar;

		m_fnCallback = [&condVar](bool)
		{
			condVar.notify_one();
		};

		while (m_status == Pending)
		{
			condVar.wait(lock);
		}
	}

	RethrowException();
	return m_status == Completed;
}

class CFuncWrapper : public CRefCounted
{
public:
	CFuncWrapper(std::function<void(bool)>&& fnCallback, bool bResult) : m_fnCallback(std::move(fnCallback)), m_bResult(bResult)
	{
	}

	~CFuncWrapper() override
	{
		// Must callback even if it was added to the TaskThread but canceled
		if (m_fnCallback)
			m_fnCallback(false);
	}

	void Execute()
	{
		if (m_fnCallback)
		{
			m_fnCallback(m_bResult);
			m_fnCallback = nullptr;
		}
	}

private:
	std::function<void(bool)> m_fnCallback;
	bool m_bResult;
};

void CAssocStateBase::Complete()
{
	std::unique_lock lock(m_mutex);

	if (m_status == Pending)
		m_status = Canceled;

	if (auto fnCallback = std::move(m_fnCallback))
	{
		ATID atid = m_atid;
		bool bResult = m_status == Completed;
		lock.unlock();

		if (atid == 0 || atid == g_asyncLib.InnerGetAtid())
			fnCallback(bResult);
		else
			g_asyncLib.InnerPostTask(atid, CAutoRefPtr(new CFuncWrapper(std::move(fnCallback), bResult)));
	}
}

void CAssocStateBase::CaptureThread()
{
	std::unique_lock lock(m_mutex);
	m_atid = g_asyncLib.InnerGetAtid();
}

void CAssocStateBase::SetException(const std::exception_ptr& exception)
{
	std::unique_lock lock(m_mutex);
	assert(m_status == Pending && exception);

	if (m_status == Pending && exception)
	{
		m_exception = exception;
		m_status = Completed;
	}
}

bool CAssocStateBase::MoveException(CAssocStateBase& stateFrom)
{
	std::scoped_lock lock(m_mutex, stateFrom.m_mutex);
	assert(m_status == Pending);

	if (m_status == Pending && stateFrom.m_exception)
	{
		stateFrom.m_bExHandled = true;
		m_exception = stateFrom.m_exception;
		m_status = Completed;
		return true;
	}

	return false;
}

void CAssocStateBase::RethrowException()
{
	if (m_exception)
	{
		m_bExHandled = true;
		std::rethrow_exception(m_exception);
	}
}

} // namespace AsyncLib
